{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/camenduru/diffusers-image-outpaint-jupyter/blob/main/diffusers_image_outpaint_jupyter.ipynb)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VjYy0F2gZIPR"
      },
      "outputs": [],
      "source": [
        "%cd /content\n",
        "!git clone -b dev https://github.com/camenduru/diffusers-image-outpaint-hf\n",
        "%cd /content/diffusers-image-outpaint-hf\n",
        "!pip install transformers accelerate diffusers\n",
        "\n",
        "!apt -y install -qq aria2\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/raw/main/lightning/model_index.json -d /content/model/lightning -o model_index.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/raw/main/lightning/scheduler/scheduler_config.json -d /content/model/lightning/scheduler -o scheduler_config.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/raw/main/lightning/text_encoder/config.json -d /content/model/lightning/text_encoder -o config.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/resolve/main/lightning/text_encoder/model.fp16.safetensors -d /content/model/lightning/text_encoder -o model.fp16.safetensors\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/raw/main/lightning/text_encoder_2/config.json -d /content/model/lightning/text_encoder_2 -o config.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/resolve/main/lightning/text_encoder_2/model.fp16.safetensors -d /content/model/lightning/text_encoder_2 -o model.fp16.safetensors\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/raw/main/lightning/tokenizer/merges.txt -d /content/model/lightning/tokenizer -o merges.txt\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/raw/main/lightning/tokenizer/special_tokens_map.json -d /content/model/lightning/tokenizer -o special_tokens_map.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/raw/main/lightning/tokenizer/tokenizer_config.json -d /content/model/lightning/tokenizer -o tokenizer_config.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/raw/main/lightning/tokenizer/vocab.json -d /content/model/lightning/tokenizer -o vocab.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/raw/main/lightning/tokenizer_2/merges.txt -d /content/model/lightning/tokenizer_2 -o merges.txt\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/raw/main/lightning/tokenizer_2/special_tokens_map.json -d /content/model/lightning/tokenizer_2 -o special_tokens_map.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/raw/main/lightning/tokenizer_2/tokenizer_config.json -d /content/model/lightning/tokenizer_2 -o tokenizer_config.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/raw/main/lightning/tokenizer_2/vocab.json -d /content/model/lightning/tokenizer_2 -o vocab.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/raw/main/lightning/unet/config.json -d /content/model/lightning/unet -o config.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/resolve/main/lightning/unet/diffusion_pytorch_model.fp16.safetensors -d /content/model/lightning/unet -o diffusion_pytorch_model.fp16.safetensors\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/raw/main/lightning/unet/diffusion_pytorch_model.safetensors.index.json -d /content/model/lightning/unet -o diffusion_pytorch_model.safetensors.index.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/raw/main/lightning/vae/config.json -d /content/model/lightning/vae -o config.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/resolve/main/lightning/vae/diffusion_pytorch_model.fp16.safetensors -d /content/model/lightning/vae -o diffusion_pytorch_model.fp16.safetensors\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/raw/main/vae-fix/config.json -d /content/model/vae-fix -o config.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/resolve/main/vae-fix/diffusion_pytorch_model.safetensors -d /content/model/vae-fix -o diffusion_pytorch_model.safetensors\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/raw/main/union/config_promax.json -d /content/model/union -o config_promax.json\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/outpaint/resolve/main/union/diffusion_pytorch_model_promax.safetensors -d /content/model/union -o diffusion_pytorch_model_promax.safetensors\n",
        "\n",
        "import torch\n",
        "from diffusers import AutoencoderKL\n",
        "from diffusers.models.model_loading_utils import load_state_dict\n",
        "from controlnet_union import ControlNetModel_Union\n",
        "from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline\n",
        "from PIL import Image, ImageDraw\n",
        "\n",
        "config = ControlNetModel_Union.load_config(\"/content/model/union/config_promax.json\")\n",
        "controlnet_model = ControlNetModel_Union.from_config(config)\n",
        "state_dict = load_state_dict(\"/content/model/union/diffusion_pytorch_model_promax.safetensors\")\n",
        "model, _, _, _, _ = ControlNetModel_Union._load_pretrained_model(controlnet_model, state_dict, \"/content/model/union/diffusion_pytorch_model_promax.safetensors\", \"/content/model/union\")\n",
        "model.to(device=\"cuda\", dtype=torch.float16)\n",
        "vae = AutoencoderKL.from_pretrained(\"/content/model/vae-fix\", torch_dtype=torch.float16).to(\"cuda\")\n",
        "pipe = StableDiffusionXLFillPipeline.from_pretrained(\"/content/model/lightning\", torch_dtype=torch.float16, vae=vae, controlnet=model, variant=\"fp16\").to(\"cuda\")\n",
        "\n",
        "def infer(image, width, height, overlap_width, num_inference_steps, prompt_input=None):\n",
        "    source = image\n",
        "    target_size = (width, height)\n",
        "    overlap = overlap_width\n",
        "\n",
        "    if source.width < target_size[0] and source.height < target_size[1]:\n",
        "        scale_factor = min(target_size[0] / source.width, target_size[1] / source.height)\n",
        "        new_width = int(source.width * scale_factor)\n",
        "        new_height = int(source.height * scale_factor)\n",
        "        source = source.resize((new_width, new_height), Image.LANCZOS)\n",
        "\n",
        "    if source.width > target_size[0] or source.height > target_size[1]:\n",
        "        scale_factor = min(target_size[0] / source.width, target_size[1] / source.height)\n",
        "        new_width = int(source.width * scale_factor)\n",
        "        new_height = int(source.height * scale_factor)\n",
        "        source = source.resize((new_width, new_height), Image.LANCZOS)\n",
        "\n",
        "    margin_x = (target_size[0] - source.width) // 2\n",
        "    margin_y = (target_size[1] - source.height) // 2\n",
        "\n",
        "    background = Image.new('RGB', target_size, (255, 255, 255))\n",
        "    background.paste(source, (margin_x, margin_y))\n",
        "\n",
        "    mask = Image.new('L', target_size, 255)\n",
        "    mask_draw = ImageDraw.Draw(mask)\n",
        "    mask_draw.rectangle([\n",
        "        (margin_x + overlap, margin_y + overlap),\n",
        "        (margin_x + source.width - overlap, margin_y + source.height - overlap)\n",
        "    ], fill=0)\n",
        "\n",
        "    cnet_image = background.copy()\n",
        "    cnet_image.paste(0, (0, 0), mask)\n",
        "\n",
        "    final_prompt = \"high quality\"\n",
        "    if prompt_input and prompt_input.strip():\n",
        "        final_prompt += \", \" + prompt_input\n",
        "\n",
        "    (\n",
        "        prompt_embeds,\n",
        "        negative_prompt_embeds,\n",
        "        pooled_prompt_embeds,\n",
        "        negative_pooled_prompt_embeds,\n",
        "    ) = pipe.encode_prompt(final_prompt, \"cuda\", True)\n",
        "\n",
        "    results = []\n",
        "    \n",
        "    for image in pipe(\n",
        "        prompt_embeds=prompt_embeds,\n",
        "        negative_prompt_embeds=negative_prompt_embeds,\n",
        "        pooled_prompt_embeds=pooled_prompt_embeds,\n",
        "        negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,\n",
        "        image=cnet_image,\n",
        "        num_inference_steps=num_inference_steps\n",
        "    ):\n",
        "        results.append((cnet_image, image))\n",
        "\n",
        "    image = image.convert(\"RGBA\")\n",
        "    cnet_image.paste(image, (0, 0), mask)\n",
        "\n",
        "    results.append((background, cnet_image))\n",
        "\n",
        "    return results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "input_image = Image.open(\"/content/diffusers-image-outpaint-hf/examples/example_3.jpg\")\n",
        "width = 1280\n",
        "height = 720\n",
        "overlap_width = 42\n",
        "num_inference_steps = 8\n",
        "prompt_input = None\n",
        "output_image = infer(input_image, height, width, overlap_width, num_inference_steps, None)\n",
        "output_image[num_inference_steps+1][1]"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
